import sys
import csv

import cvxpy as cp
import numpy as np
import lib.tools as my_tools

from tqdm import tqdm
from qpsolvers import solve_qp
from lib.model import system_model
from itertools import product
from loguru import logger

interception_system = system_model("interception_system")
logger.add("solve.log")

def solve_dataset_n3(yaml_name):
    # 加载参数
    config = my_tools.load_yaml(yaml_name)
    logger.info(f"Params {config} have been loaded.")
    eta = config['eta']
    w_u = config['w_u']
    w_omega = config['w_omega']
    
    # # 控制量
    # omega_x = np.linspace(-0.1, 0.1, num=int((0.1 - (-0.1)) / 0.005)+1, dtype=np.float64)
    # omega_y = np.linspace(-0.1, 0.1, num=int((0.1 - (-0.1)) / 0.005)+1, dtype=np.float64)

    # data = [["omega_x len", len(omega_x)],
    #         ["omega_y len", len(omega_y)]]
    # my_tools.print_table(data, "Control seed lengths")
    
    # 状态量
    omega_z = np.linspace(-0.8, 0.8, num=int((0.8 - (-0.8)) / 0.1)+1, dtype=np.float64)
    gen_n3d = interception_system.generate_n3(bound = 0.1)
    n3d = list(gen_n3d)
    gen_n3 = interception_system.generate_n3(bound = 0.1)
    n3 = list(gen_n3)

    data = [
        ["omega_z", len(omega_z)],
        ["n3", len(n3)],
        ["n3d", len(n3d)],
        ["all combinations", len(omega_z) * len(n3) * len(n3d)]
    ]
    my_tools.print_table(data, "State seed lengths")

    # 根据状态组合数量添加带进度条的循环
    state_pdt = product(range(len(omega_z)), range(len(n3)), range(len(n3d)))
    state_len = len(omega_z) * len(n3) * len(n3d)
    state_tqdm = tqdm(state_pdt, total = state_len, ncols=110)

    with open('solved_dataset_attitude.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['D_value', 'omega_x', 'omega_y', 'omega_z', 'n3_x', 'n3_y', 'n3_z', 'n3_xd', 'n3_yd', 'n3_zd', 'desired_D_value','opt_error'])
        data_item = []
        for s1, s2, s3 in state_tqdm:
            if np.array_equal(n3[s2], n3d[s3]):
                data_item.append([0, 0, 0, omega_z[s1], n3[s2][0], n3[s2][1], n3[s2][2], n3d[s3][0], n3d[s3][1], n3d[s3][2], 0, 0])
                writer.writerow(data_item[-1])
            else:
                tilde_n3x = n3[s2][0] - n3d[s3][0]
                tilde_n3y = n3[s2][1] - n3d[s3][1]
                tilde_n3z = n3[s2][2] - n3d[s3][2]                
                N =  np.array([[tilde_n3y*n3[s2][2] - tilde_n3z*n3[s2][1], tilde_n3z*n3[s2][0] - tilde_n3x*n3[s2][2]]])
                P = N.T @ (w_omega*N) + w_u*np.eye(2)
                n_c = -(tilde_n3x*n3[s2][1]*omega_z[s1] - tilde_n3y*n3[s2][0]*omega_z[s1] + eta*(tilde_n3x**2 + tilde_n3y**2))
                q = -n_c*w_omega*N.T
                lb = np.array([-2, -2])
                ub = np.array([2, 2])
                omega2 = solve_qp(P, q, lb=lb, ub=ub, solver="cvxopt")

                # 其他辅助数值计算
                D_value = interception_system.cal_D_value_n3(n_3 = n3[s2], 
                                                            n_3d = n3d[s3], 
                                                            omega = np.array([omega2[0], omega2[1], omega_z[s1]])
                                                            )
                error = N @ omega2.reshape(2,1) - n_c
                desired_D_value = -eta*(tilde_n3x**2 + tilde_n3y**2)
                data_item.append([D_value[0,0], omega2[0], omega2[1], omega_z[s1], n3[s2][0], n3[s2][1], n3[s2][2], n3d[s3][0], n3d[s3][1], n3d[s3][2], desired_D_value, error.item()])
                writer.writerow(data_item[-1])
                # break
        my_tools.save_dataset(data_item, 'sorted_dataset_attitude.pickle')

COMMANDS = {
    "attitude": "Sort attitude control dataset"
}

if __name__ == "__main__":
    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        # param miss
        logger.error("YAML file name miss.")

    if len(sys.argv) == 3:
        if sys.argv[1] == 'attitude':
            yaml_name = sys.argv[2]
            solve_dataset_n3(yaml_name)
        else:
            my_tools.print_help(sys.argv[0], COMMANDS)